#include <vector>
#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include <iostream>
#include <cstdlib>
#include <time.h>

#define protected public
#define private public

#undef private
#undef protected

#include "acl_rt.h"
#include "runtime/base.h"
#include "runtime/mem.h"

using namespace std;
using namespace testing;

extern "C" aclError GetAlignedSize(const size_t size, size_t *alignedSize);

class AclTest : public testing::Test {
protected:
  void SetUp(){}
  void TearDown(){}
};

TEST_F(AclTest, memory_malloc_device) {
  void *devPtr = NULL;
  size_t size = 1;

  aclError ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_HUGE_ONLY);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_NORMAL_ONLY);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_HUGE_FIRST_P2P);
  EXPECT_EQ(ret, ACL_ERROR_FEATURE_UNSUPPORTED);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_HUGE_ONLY_P2P);
  EXPECT_EQ(ret, ACL_ERROR_FEATURE_UNSUPPORTED);

  ret = aclrtMalloc(&devPtr, size, (aclrtMemMallocPolicy)(ACL_MEM_TYPE_HIGH_BAND_WIDTH | ACL_MEM_MALLOC_NORMAL_ONLY_P2P));
  EXPECT_EQ(ret, ACL_ERROR_FEATURE_UNSUPPORTED);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_TYPE_HIGH_BAND_WIDTH);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, ACL_MEM_TYPE_LOW_BAND_WIDTH);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, (aclrtMemMallocPolicy)(ACL_MEM_TYPE_HIGH_BAND_WIDTH | ACL_MEM_TYPE_LOW_BAND_WIDTH));
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMalloc(&devPtr, size, (aclrtMemMallocPolicy)6);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  ret = aclrtMalloc(NULL, size, ACL_MEM_MALLOC_HUGE_FIRST);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  size = (size_t)(-1);
  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_NORMAL_ONLY);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  size = (size_t)(0);
  ret = aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_NORMAL_ONLY);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
}

TEST_F(AclTest, memory_free_device) {
  void *devPtr = (void *)0x01;
  aclError ret = aclrtFree(NULL);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
  ret = aclrtMalloc(&devPtr, 1, ACL_MEM_MALLOC_HUGE_ONLY);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtFree(devPtr);
  EXPECT_EQ(ret, ACL_SUCCESS);
}

TEST_F(AclTest, memory_memset) {
    void *ptr = (void *)0x01;
    aclError ret = aclrtMemset(ptr, 1, 1, 1);
    EXPECT_EQ(ret, ACL_SUCCESS);

    ret = aclrtMemset(NULL, 1, 1, 1);
    EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

    ret = aclrtMemset(ptr, 1, 1, 2);
    EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
}

TEST_F(AclTest, memory_memcpy) {
  void *dst = (void *)0x01;
  void *src = (void *)0x02;
  aclError ret = aclrtMemcpy(dst, 1, src, 1, ACL_MEMCPY_HOST_TO_HOST);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMemcpy(dst, 1, src, 1, ACL_MEMCPY_HOST_TO_DEVICE);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMemcpy(dst, 1, src, 1, ACL_MEMCPY_DEVICE_TO_DEVICE);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMemcpy(dst, 1, src, 1, ACL_MEMCPY_DEVICE_TO_HOST);
  EXPECT_EQ(ret, ACL_SUCCESS);

  ret = aclrtMemcpy(NULL, 1, src, 1, ACL_MEMCPY_HOST_TO_HOST);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  ret = aclrtMemcpy(dst, 1, NULL, 1, ACL_MEMCPY_HOST_TO_HOST);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  ret = aclrtMemcpy(dst, 1, src, 1, (aclrtMemcpyKind)0x7FFFFFFFF);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);

  ret = aclrtMemcpy(dst, 1, src, 2, ACL_MEMCPY_HOST_TO_HOST);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
}

TEST_F(AclTest, memory_getMemInfo) {
  size_t free;
  size_t total;

  aclError ret = aclrtGetMemInfo(ACL_DDR_MEM_HUGE, &free, &total);
  EXPECT_EQ(ret, ACL_SUCCESS);
  ret = aclrtGetMemInfo(ACL_DDR_MEM_HUGE, &free, NULL);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
  ret = aclrtGetMemInfo(ACL_DDR_MEM_HUGE, NULL, &total);
  EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
}

TEST_F(AclTest, GetAlignedSize) {
  size_t alignedSize;
  size_t size;
  srand((unsigned int)time(NULL));
  for(int i = 0; i < 50; i++) {
    size = (rand() % (2000 - 1 + 1)) + 1;
    if (GetAlignedSize(size, &alignedSize) == ACL_SUCCESS) {
      EXPECT_EQ((alignedSize - 32) % 32, 0);
    }
  }
}